import torch
from torch import nn

# ein 就是爱因斯坦的ein，sum就是求和。einsum就是爱因斯坦求和约定，其实作用就是把求和(乘法)符号省略

# 求和
x = torch.randn((3, 4, 5))
y = torch.einsum("ijk -> ij", x)
y = torch.einsum("ijk -> ", x)

# 转置
x = torch.randn((3, 4, 5))
y = torch.einsum("ijk -> jik", x)
y = torch.einsum("ijk -> kij", x)

# 张量乘法
x = torch.randn((3, 4))
y = torch.randn((4, 5))
z = torch.einsum("i j,j k -> i k", x, y)

x = torch.randn((3, 4, 5, 6))
y = torch.randn((3, 4, 8, 9))
z = torch.einsum("i j k l,i j m n -> k l m n", x, y)

# 张量乘法的规则还未明晰，只知道它需要一些相同的维度，并在做完乘法以后可以替换之
